import time
import threading
import numpy as np
import pandas as pd
import sqlite3
from tqsdk import TqApi, TqAuth, TqKq
from MyTT import EMA
from rich.console import Console
from rich.theme import Theme

# 定义颜色主题
custom_theme = Theme({
    "info": "cyan",
    "warning": "yellow",
    "error": "red",
    "success": "green",
    "trade": "bold magenta"
})
console = Console(theme=custom_theme)

class ThreadedTradingBot:
    def __init__(self):
        self.api = None
        self.running = True
        self.threads = []
        self.lock = threading.Lock()
        self.console = console

    def get_monitored_symbols(self):
        """从数据库获取监控品种"""
        conn = sqlite3.connect('trade_monitor.db')
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM products WHERE is_monitored=1")
        products = cursor.fetchall()
        conn.close()
        return [row[2] for row in products]

    def run(self):
        """主运行函数"""
        try:
            self.api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba999"))
            symbols = self.get_monitored_symbols()
            self.console.print(f"开始监控品种: {symbols}", style="info")
            
            # 为每个品种创建交易线程
            for symbol in symbols:
                thread = threading.Thread(
                    target=self.trade_loop,
                    args=(symbol,),
                    daemon=True
                )
                thread.start()
                self.threads.append(thread)
            
            # 主线程保持运行
            while self.running:
                time.sleep(1)
                
        except Exception as e:
            self.console.print(f"程序出错: {e}", style="error")
        finally:
            self.running = False
            for thread in self.threads:
                thread.join()
            if self.api:
                self.api.close()

    def trade_loop(self, symbol):
        """单个品种的交易循环"""
        try:
            quote = self.api.get_quote(symbol)
            klines = self.api.get_kline_serial(symbol, 60, 500)

            # 注册更新通知
            update_event = threading.Event()
            self.api.register_update_notify(update_event.set)
            
            while self.running:
                # 等待更新事件或超时
                update_event.wait(timeout=1)
                update_event.clear()

                if self.api.is_changing(quote):
                    self.check_stoploss(symbol, quote, klines)
                    
        except Exception as e:
            self.console.print(f"处理品种 {symbol} 时出错: {e}", style="error")

    def check_stoploss(self, symbol, quote, klines):
        """止损逻辑"""
        self.console.print(f"{time.strftime('%X')} {symbol} 价格: {quote.last_price}", style="info")
        
        try:
            # 确保 klines 的各列是可订阅的 Series 对象
            if hasattr(klines, 'close') and hasattr(klines, 'low') and hasattr(klines, 'open') and hasattr(klines, 'high'):
                q = (3 * klines.close + klines.low + klines.open + klines.high) / 6
            else:
                self.console.print(f"{symbol} 的 K 线数据缺少必要列，无法计算技术指标", style="error")
                return
            # 后续可以继续添加计算逻辑，如果有使用 q 的地方，确保使用正确
        except Exception as e:
            self.console.print(f"计算 {symbol} 技术指标时出错: {e}", style="error")
            return

        weights = np.arange(26, 0, -1)
        ma = []
        for i in range(len(q)):
            if i < 25:
                ma.append(np.nan)
                continue
            try:
                window = q.iloc[i - 25:i + 1]
                ma.append(np.dot(window, weights) / 351)
            except Exception as e:
                self.console.print(f"计算 {symbol} 的移动平均线时出错: {e}", style="error")
                ma.append(np.nan)

        try:
            ema_7 = EMA(pd.Series(ma), 7).tolist()
        except Exception as e:
            self.console.print(f"计算 {symbol} 的 EMA7 时出错: {e}", style="error")
            return

        # 使用锁保证线程安全
        with self.lock:
            position = self.api.get_position(symbol)
            
            # 平仓条件判断
            if len(ma) >= 2 and len(ema_7) >= 2:
                if position.pos_short > 0 and ma[-2] > ema_7[-2]:
                    self.api.insert_order(
                        symbol=symbol,
                        direction="BUY",
                        offset="CLOSE",
                        volume=position.pos_short
                    )
                    self.console.print(
                        f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}", 
                        style="trade"
                    )
                    
                if position.pos_long > 0 and ma[-2] < ema_7[-2]:
                    self.api.insert_order(
                        symbol=symbol,
                        direction="SELL",
                        offset="CLOSE",
                        volume=position.pos_long
                    )
                    self.console.print(
                        f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}", 
                        style="trade"
                    )

if __name__ == "__main__":
    bot = ThreadedTradingBot()
    try:
        bot.run()
    except KeyboardInterrupt:
        bot.running = False
        console.print("\n程序已停止", style="warning")